from loadDataset import LoadBCICompDataSet
import numpy as np
import os
from tqdm import tqdm


def extractDataAndSave(matFilePath: str, saveRootPath: str) -> None:
    '''
    Description:
        这个函数是用来将mat文件中是每个Target的每个Channel的信号进行分类，并保存到文件中

        文件的顺序如下
            └─Subject
                └─TargetNum
                    └─ChannelNum


    Parameters:
        matFilePath:str
            mat文件所在的路径
        saveRootPath:str
            提取信号所在的根文件夹
    '''
    dataSet = LoadBCICompDataSet.LoadBCICompDataSet(matFilePath)
    signal = dataSet["Signal"]
    flash = dataSet["Flashing"]
    code = dataSet["StimulusCode"]
    saveData = np.zeros((85, 64, 12, 15, 240))
    windows = 240
    for i in tqdm(range(0, np.size(signal, 0)), desc='Extracting'):
        testSignal: np.array = signal[i]
        testFlash: np.array = flash[i]
        testCode: np.array = code[i]
        for j in range(0, np.size(testSignal, 1)):
            # print(test[:, i].size)
            currentChannel = testSignal[:, j]
            # print(type(currentChannel.size)) # int
            rowColCount = np.zeros(12)
            for k in range(1, testFlash.size):
                if testFlash[k] == 0 and testFlash[k - 1] == 1:
                    rowCol = int(testCode[k - 1]) - 1
                    saveData[i, j, rowCol, int(rowColCount[rowCol]), :] = currentChannel[
                                                                          k - 24:k + windows - 24]  # 按照案例给的
    # print(saveData[0, 0, 0, 0, :])
    root = saveRootPath
    # subPath = os.path.join(root, "subjectA")
    # if not os.path.exists(subPath):
    #     os.mkdir(subPath)
    if not os.path.exists(root):
        os.mkdir(root)
    os.chdir(root)
    for i in tqdm(range(0, np.size(saveData, 0))):
        targetPath = os.path.join(root, "Target_" + str(i + 1))
        if not os.path.exists(targetPath):
            os.mkdir(targetPath)
        os.chdir(targetPath)
        for j in range(0, np.size(saveData, 1)):
            a = saveData[i, j, :, :]
            # print(a.shape)
            np.save("channel_" + str(j + 1) + ".npy", a)


def extractData(matFilePath: str) -> np.ndarray:
    '''
    Description:
        这个函数是用来读取提取信号保存到文件中的numpy数组

    Parameters:
        matFilePath:str
            提取信号后的根路径

    Returns:
        saveData:np.ndarray
            返回提取信号数组，大小是[85, 64, 12, 15, 240]
    '''
    dataSet = LoadBCICompDataSet.LoadBCICompDataSet(matFilePath)
    signal = dataSet["Signal"]
    flash = dataSet["Flashing"]
    code = dataSet["StimulusCode"]
    saveData = np.zeros((85, 64, 12, 15, 240))
    windows = 240
    for i in tqdm(range(0, np.size(signal, 0)), desc="Extracting"):
        testSignal: np.array = signal[i]
        testFlash: np.array = flash[i]
        testCode: np.array = code[i]
        for j in range(0, np.size(testSignal, 1)):
            # print(test[:, i].size)
            currentChannel = testSignal[:, j]
            # print(type(currentChannel.size)) # int
            rowColCount = np.zeros(12)
            for k in range(1, testFlash.size):
                if testFlash[k] == 0 and testFlash[k - 1] == 1:
                    rowCol = int(testCode[k - 1]) - 1
                    saveData[i, j, rowCol, int(rowColCount[rowCol]), :] = currentChannel[
                                                                          k - 24:k + windows - 24]  # 按照案例给的
    return saveData


def LoadSignalData(subjectRootPath: str, convertToF: bool = False) -> np.ndarray:
    '''
    Description:
        这个函数是用来读取提取信号保存到文件中的numpy数组

    Parameters:
        subjectRootPath:str
            提取信号后的根路径
        convertToF: bool
            是否转换为Float32类型

    Returns:
        datas:Tuplenp.ndarray
            返回信号数组，大小是[85, 64, 12, 15, 240]
    '''
    datas = np.zeros((85, 64, 12, 15, 240))
    paths = os.listdir(subjectRootPath)
    countTarget = 0
    for i in tqdm(paths):
        files = os.path.join(subjectRootPath, i)
        countChannel = 0
        for j in files:
            file = os.path.join(files, j)
            data = np.load(file)
            datas[countTarget, countChannel, :, :, :] = data
            countChannel += 1
        countTarget += 1
    if convertToF:
        datas = datas.astype(np.float32)
    return datas


if __name__ == '__main__':
    extractDataAndSave(r"F:\DataSet\BCICompDIII\DataSet\BCI_Comp_III_Wads_2004\Subject_A_Train.mat",
                       r"F:\DataSet\BCICompDIII\DataSet\MyExtract\SubjectA")
    extractDataAndSave(r"F:\DataSet\BCICompDIII\DataSet\BCI_Comp_III_Wads_2004\Subject_B_Train.mat",
                       r"F:\DataSet\BCICompDIII\DataSet\MyExtract\SubjectB")
